import wx
import os
import numpy
import time
from threading import Thread

from ViewPanel import ViewPanel
from Buffers import buffers
from NotifyEvent import NotifyEvent

threshold_buffers = {}

from ConvolutionFrame import ConvolutionFrame
from MorphologyFrame import MorphologyFrame

import psyco
psyco.full()

class MainFrame(wx.Frame):
	def __init__(self, parent, ID, title):
		wx.Frame.__init__(self, parent, wx.ID_ANY, title, size=(1024,768))

		self.filename = ""
		
		self.initStatusBar()
		self.createMenuBar()
		self.createToolBar()
		self.viewpanel = ViewPanel(self, wx.ID_ANY)

	def initStatusBar(self):
		self.statusbar = self.CreateStatusBar()
		self.statusbar.SetFieldsCount(3)
		self.statusbar.SetStatusWidths([-1, -2, -3])

	def createStatusBar(self):
		self.statusbar = wx.StatusBar(self, style=wx.SB_FLAT)
		self.statusbar.SetFieldsCount(3)
		self.statusbar.SetStatusWidths([-1, -1, -4])
		self.SetStatusBar(self.statusbar)

	def menuData(self):
		return [("&File", (
					("&New", "New Sketch file", self.OnNew),
					("&Open", "Open sketch file", self.OnOpen),
					("&Save", "Save sketch file", self.OnSave),
					("", "", ""),
					("&History", "History of sketch file", self.OnBuffers),
					("", "", ""),
					("&Exit", "Exit", self.OnCloseWindow))),
				("&Help", (
					("&New", "New Sketch file", self.OnNew),
					("", "", ""),
					("&About", "About", self.OnAbout)))]

	def createMenuBar(self):
		menuBar = wx.MenuBar()
		for eachMenuData in self.menuData():
			menuLabel = eachMenuData[0]
			menuItems = eachMenuData[1]
			menuBar.Append(self.createMenu(menuItems), menuLabel)
		self.SetMenuBar(menuBar)
		
	def createMenu(self, menuData):
		menu = wx.Menu()
		for eachItem in menuData:
			if len(eachItem) == 2:
				label = eachItem[0]
				subMenu = self.createMenu(eachItem[1])
				menu.AppendMenu(wx.NewId(), label, subMenu)
			else:
				self.createMenuItem(menu, *eachItem)
		return menu
	
	def createMenuItem(self, menu, label, status, handler, kind=wx.ITEM_NORMAL):
		if not label:
			menu.AppendSeparator()
			return
		menuItem = menu.Append(-1, label, status, kind)
		self.Bind(wx.EVT_MENU, handler, menuItem)

	def createToolBar(self):
		toolbar = self.CreateToolBar()
		for each in self.toolbarData():
			self.createSimpleTool(toolbar, *each)
		toolbar.Realize()
		self.SetToolBar(toolbar)

	def createSimpleTool(self, toolbar, label, filename, handler):
		if not label:
			toolbar.AddSeparator()
			return
		bmp = wx.Image(filename, wx.BITMAP_TYPE_ANY).ConvertToBitmap()
		tool = toolbar.AddSimpleTool(wx.ID_ANY, bmp, label)
		self.Bind(wx.EVT_MENU, handler, tool)

	def toolbarData(self):
		return (
			    ("Open", os.path.join("icon", "open.png"), self.OnOpen),
			    ("Save", os.path.join("icon", "save.png"), self.OnSave),
				("History", os.path.join("icon", "history.png"), self.OnBuffers),
				("Remove", os.path.join("icon", "remove.png"), self.OnRemove),
				("", "", ""),
				("Histogram", os.path.join("icon", "histogram.png"), self.OnHistogram),
				("Convolution", os.path.join("icon", "convolution.png"), self.OnConvolution),
				("Mathematical Morphology", os.path.join("icon", "morphology.png"), self.OnMorphology)
		)

	def OnCloseWindow(self, event):
		self.Destroy()
	
	def OnBuffers(self, event):
		menu = wx.Menu()
		for id, name in buffers.id_name_pairs():
			menu.Append(id, os.path.basename(name))
		menu.Bind(wx.EVT_MENU, self.OnSelectBuffer)
		self.PopupMenu(menu, (0, 0))

	def OnNew(self, event): pass
	
	wildcard = "JPEG files (*.jpg)|*.jpg|"\
	            "GIF files (*.gif)|*.gif|"\
				"BMP files (*.bmp)|*.bmp|" \
				"PNG files (*.png)|*.png|" \
				"ALL files (*.*)|*.*|"

	def OnOpen(self, event):
		dlg = wx.FileDialog(self, "Open image file...", os.getcwd(),
						style=wx.OPEN, wildcard=self.wildcard)
		if dlg.ShowModal() == wx.ID_OK:
			filename = dlg.GetPath()
			if not buffers.has_name(filename):
				self.ReadFile(filename)
		dlg.Destroy()

	def OnRemove(self, event):
		name = self.viewpanel.GetImageName()
		buffers.remove(name)
		self.viewpanel.ClearImage()

	def OnSave(self, event):
		if not self.filename:
			self.OnSaveAs(event)
		else:
			self.SaveFile()

	def OnHistogram(self, event):
		source_name = self.viewpanel.GetPreviousImageName()
		if not source_name:
			return
		source = buffers.get_image_by_name(source_name)
		image = self.generateHistogram(source)

		new_name = source_name + " - histogram"
		name = buffers.add(new_name, image, source_name)
		self.displayBuffer(name)

		# generate threshold
		histogram_view = self.viewpanel.GetViewByName(new_name)
		thresholding_thread = Thresholding(source_name, histogram_view)
		thresholding_thread.run()
		#thresholding_thread.start()

	def OnConvolution(self, event):
		try:
			self.convolution_frame
		except AttributeError:
			self.convolution_frame = ConvolutionFrame(self, wx.ID_ANY, 'convolution', (50, 255))
		self.convolution_frame.Show()

		source_name = self.viewpanel.GetPreviousImageName()
		if not source_name:
			return
		source = buffers.get_image_by_name(source_name)
		self.convolution_frame.set_source(source_name, source)

	def OnMorphology(self, event):
		try:
			self.morphology_frame
		except AttributeError:
			self.morphology_frame = MorphologyFrame(self, wx.ID_ANY, 'morphology', (50, 305))
		self.morphology_frame.Show()

		source_name = self.viewpanel.GetPreviousImageName()
		if not source_name:
			return
		source = buffers.get_image_by_name(source_name)
		self.morphology_frame.set_source(source_name, source)

	def OnSelectBuffer(self, event):
		name = buffers.get_name_by_id(event.GetId())
		self.displayBuffer(name)

	def displayBuffer(self, name):
		image = buffers.get_image_by_name(name)
		self.viewpanel.SetImage(name, image)

	def ReadFile(self, filename):
		try:
			image = wx.Image(filename, type=wx.BITMAP_TYPE_ANY)
		except:
			wx.MessageBox("oops! file type not supported.",
							style=wx.OK|wx.ICON_EXCLAMATION)
		name = buffers.add(filename, image)
		self.displayBuffer(name)

	def SaveFile(self):
		if self.filename:
			data = self.viewpanel.GetLinesData()
			f = open(self.filename, 'w')
			cPickle.dump(data, f)
			f.close()
				
	def OnSaveAs(self, event):
		dlg = wx.FileDialog(self, "Save img as...", os.getcwd(),
						style=wx.SAVE | wx.OVERWRITE_PROMPT,
						wildcard = self.wildcard)
		if dlg.ShowModal() == wx.ID_OK:
			filename = dlg.GetPath()
			if not os.path.splitext(filename)[1]:
				filename = filename + '.img'
			self.filename = filename
			self.SaveFile()
			self.SetTitle(self.title + ' -- ' + self.filename)
		dlg.Destroy()
	def OnAbout(self,event):
		dlg = wx.MessageDialog(self, 'message','MessageDialog', wx.OK | wx.ICON_INFORMATION) 
		dlg.ShowModal()
		dlg.Destroy()

	def generateHistogram(self, source):
		histogram = numpy.zeros((4,256), dtype=int)
		w, h = (source.GetWidth(), source.GetHeight())
		for i in range(w):
			for j in range(h):
				histogram[0, source.GetRed(i, j)] += 1
				histogram[1, source.GetGreen(i, j)] += 1
				histogram[2, source.GetBlue(i, j)] += 1
		histogram[3] = histogram[0] + histogram[1] + histogram[2]

		histogram_width = 256*2
		histogram_height = 100
		histogram_vertical_gap = 5

		scale = []
		scale.append(float(histogram_height) / histogram[0].max())
		scale.append(float(histogram_height) / histogram[1].max())
		scale.append(float(histogram_height) / histogram[2].max())
		scale.append(float(histogram_height) / histogram[3].max())

		bitmap = wx.EmptyBitmap(histogram_width, (histogram_height+histogram_vertical_gap)*4)
		dc = wx.MemoryDC(bitmap)
		dc.SetBackground(wx.Brush('black'))
		dc.Clear()

		colours = ['red', 'green', 'blue', 'white']
		i = 0
		start_y =  histogram_height - 1
		for colour in colours:
			dc.SetPen(wx.Pen(colour, 1))
			start_x = 0
			for val in histogram[i]:
				diff = int(val*scale[i])
				dc.DrawLine(start_x, start_y, start_x, start_y-diff)
				start_x += 2

			y_diff = histogram_height + histogram_vertical_gap
			start_y += y_diff
			i += 1

		image = bitmap.ConvertToImage()

		return image

	def threshold_control(self, source_name, threshold, display_id):
		name = buffers.get_origin_name(source_name)
		new_name = name + " - threshold"
		dest_image = threshold_buffers[name][threshold]
		self.viewpanel.SetImage(new_name, dest_image, display_id)


class Thresholding:
	def __init__(self, source_name, view):
		#Thread.__init__(self)
		self.source_name = source_name
		self.source = buffers.get_image_by_name(self.source_name)
		self.view = view
		threshold_buffers[self.source_name] = []

	def run(self):
		for t in range(256):
			w, h = (self.source.GetWidth(), self.source.GetHeight())
			bitmap = wx.EmptyBitmap(w, h)
			dc = wx.MemoryDC(bitmap)
			dc.SetBackground(wx.Brush('black'))
			dc.Clear()
			dc.SetPen(wx.Pen('white', 1))
			for i in range(w):
				for j in range(h):
					gray = (self.source.GetRed(i, j)+self.source.GetGreen(i, j)+self.source.GetBlue(i, j))/3
					if gray >= t:
						dc.DrawPoint(i, j)
			image = bitmap.ConvertToImage()
			threshold_buffers[self.source_name].append(image)
			self.view.notify(t)

class ImageEditApp(wx.App):
	def OnInit(self):
		frame = MainFrame(None, wx.ID_ANY, title="ImageEdit")
		frame.Show()
		return True

if __name__ == '__main__':

	# Import Psyco if available
	try:
		import psyco
		psyco.full()
	except ImportError:
		print 'psyco import error'

	app = ImageEditApp(False)
	app.MainLoop()
